clear, clc;
%% 系统参数配置
Rr = 0.0013;    Xm = 2.6;       Xrr = Xm;     
Ts = 0.01;      C = 13.29;      Tf = 3.4;      
D = 5;          J = 6.5;        R = 0.0013;    
Xs = 2.9;       omega_s = 1.0;  X_sigma_r = 2.9;  
X_dr = X_sigma_r + Xm; X_qr = X_sigma_r + Xm;

%% 模拟数据生成
numTimeSteps = 2000; % 生成更长时间序列供训练
t = (0:numTimeSteps-1)*Ts;

% 输入信号生成
U_qs = 1.0 * sin(2*pi*50*t); 
U_ds = 0.5 * cos(2*pi*50*t);
U_dr = 0.2 * sin(2*pi*50*t + pi/3);
U_qr = 0.1 * cos(2*pi*50*t);
T_m = 0.8 + 0.1*sin(2*pi*50*t);
T_e = 0.6 + 0.05*cos(2*pi*50*t);
I_dl = 0.5 * sin(2*pi*0.1*t);

% 真实状态生成
n = 6; m = 6;
x = zeros(n, numTimeSteps);
z = zeros(m, numTimeSteps);

for i = 2:numTimeSteps
    x(1,i) = x(1,i-1) + Ts*((Rr/Xrr)*(U_qs(i-1)-x(2,i-1))) - x(6,i-1)*x(1,i-1);
    x(2,i) = x(2,i-1) + Ts*((Rr/Xrr)*(U_ds(i-1)-x(1,i-1))) + x(6,i-1)*x(2,i-1);
    x(3,i) = x(3,i-1) + Ts*(C/Tf*(x(4,i-1)-I_dl(i-1)));
    x(4,i) = x(4,i-1) + Ts*((1/X_dr)*(U_dr(i-1)-Rr*x(4,i-1)));
    x(5,i) = x(5,i-1) + Ts*((1/X_qr)*(U_qr(i-1)-Rr*x(5,i-1)));
    x(6,i) = x(6,i-1) + Ts*((T_m(i-1)-T_e(i-1)-D*x(6,i-1))/J);
    
    m_d = 0.8*sin(2*pi*50*t(i)); 
    m_q = 0.5*cos(2*pi*50*t(i));
    U_d = m_d*x(3,i)/2; U_q = m_q*x(3,i)/2;
    
    z(1,i) = (omega_s*Xs)/(R+Xs)*x(2,i) + (R/(R+Xs))*U_q;
    z(2,i) = (omega_s*Xs)/(R+Xs)*x(1,i) - (R/(R+Xs))*U_d;
    z(3,i) = x(3,i);
    z(4,i) = (Xs/(R+Xs))*x(3,i) - (R/(R+Xs))*x(1,i) + (omega_s*Xs)/(R+Xs)*z(1,i);
    z(5,i) = (R/(R+Xs))*x(3,i) + (Xs/(R+Xs))*x(1,i) + (omega_s*R)/(R+Xs)*z(2,i);
    z(6,i) = x(6,i);
end

%% LSTM网络训练
lookback = 5; % 使用5个历史时间步
numFeatures = m + n; % 观测+状态
numResponses = n;

% 构建训练数据集
XTrain = [];
YTrain = [];
for i = lookback+1:numTimeSteps
    input_seq = [z(:,i-lookback:i-1); x(:,i-lookback:i-1)];
    target = x(:,i);
    XTrain = cat(4, XTrain, input_seq);
    YTrain = cat(4, YTrain, target);
end

% 网络结构
layers = [ ...
    sequenceInputLayer(numFeatures)
    % sequenceInputLayer(numFeatures*lookback)
    lstmLayer(256, 'OutputMode','sequence')
    dropoutLayer(0.3)
    fullyConnectedLayer(128)
    reluLayer
    fullyConnectedLayer(numResponses)
    regressionLayer];

options = trainingOptions('adam', ...
    'MaxEpochs', 100, ...
    'MiniBatchSize', 128, ...
    'InitialLearnRate', 0.001, ...
    'LearnRateSchedule','piecewise', ...
    'LearnRateDropFactor',0.5, ...
    'LearnRateDropPeriod',50, ...
    'Shuffle','every-epoch', ...
    'Plots','training-progress', ...
    'Verbose',0);

net = trainNetwork(XTrain, YTrain, layers, options);

%% UKF-LSTM联合估计
X_est = zeros(n, numTimeSteps);
Pplus = eye(n)*0.1;
Q = diag([0.01 0.01 0.1 0.01 0.01 0.1]);
R_ukf = eye(m)*0.01;

for i = 2:numTimeSteps
    % UKF预测
    [Xminus, Pminus] = ukf_predict(X_est(:,i-1), Pplus, Q, Ts, Rr, Xrr, U_qs(i-1), U_ds(i-1),...
        U_dr(i-1), U_qr(i-1), T_m(i-1), T_e(i-1), D, J, C, Tf, I_dl(i-1));
    
    % UKF更新
    [X_est(:,i), Pplus] = ukf_update(Xminus, Pminus, z(:,i), R_ukf, U_qs(i), U_ds(i),...
        Xs, R, omega_s);
    
    % LSTM误差补偿
    if i > lookback
        lstm_input = [z(:,i-lookback:i); X_est(:,i-lookback:i)];
        lstm_input = reshape(lstm_input, [], 1);
        lstm_out = predict(net, reshape(lstm_input, numFeatures*lookback,1,1,1));
        X_est(:,i) = 0.7*X_est(:,i) + 0.3*lstm_out; % 融合比例可调
    end
    
    % 自适应Q调整
    err = z(:,i) - obs_model(X_est(:,i), U_qs(i), U_ds(i), Xs, R, omega_s);
    Q = Q * (0.95 + 0.05*norm(err));
end

%% 评估
RMSE = sqrt(mean((X_est - x).^2,2));
fprintf('状态变量RMSE:\n');
disp(RMSE');

%% UKF子函数
function [Xminus, Pminus] = ukf_predict(X, P, Q, Ts, Rr, Xrr, U_qs, U_ds, U_dr, U_qr, T_m, T_e, D, J, C, Tf, I_dl)
    n = length(X);
    alpha = 1e-3; beta = 2; kappa = 0;
    lambda = alpha^2*(n + kappa) - n;
    
    % Sigma点生成
    sigma = zeros(n, 2*n+1);
    sigma(:,1) = X;
    sqrtP = chol((n + lambda)*P);
    for i = 1:n
        sigma(:,i+1) = X + sqrtP(:,i);
        sigma(:,i+n+1) = X - sqrtP(:,i);
    end
    
    % 状态传播
    X_sig = zeros(n, 2*n+1);
    for j = 1:2*n+1
        X_sig(:,j) = state_model(sigma(:,j), Ts, Rr, Xrr, U_qs, U_ds, U_dr, U_qr, T_m, T_e, D, J, C, Tf, I_dl);
    end
    
    % 计算均值和协方差
    Wm = [lambda/(n+lambda), 0.5/(n+lambda)*ones(1,2*n)];
    Wc = Wm; Wc(1) = Wc(1) + (1 - alpha^2 + beta);
    Xminus = X_sig*Wm';
    Pminus = Q;
    for j = 1:2*n+1
        Pminus = Pminus + Wc(j)*(X_sig(:,j) - Xminus)*(X_sig(:,j) - Xminus)';
    end
end

function [X, P] = ukf_update(Xminus, Pminus, z, R, U_qs, U_ds, Xs, R_sys, omega_s)
    n = length(Xminus); m = length(z);
    alpha = 1e-3; beta = 2; kappa = 0;
    lambda = alpha^2*(n + kappa) - n;
    
    % Sigma点生成
    sigma = zeros(n, 2*n+1);
    sigma(:,1) = Xminus;
    sqrtP = chol((n + lambda)*Pminus);
    for i = 1:n
        sigma(:,i+1) = Xminus + sqrtP(:,i);
        sigma(:,i+n+1) = Xminus - sqrtP(:,i);
    end
    
    % 观测预测
    Z_sig = zeros(m, 2*n+1);
    for j = 1:2*n+1
        Z_sig(:,j) = obs_model(sigma(:,j), U_qs, U_ds, Xs, R_sys, omega_s);
    end
    
    % 计算统计量
    Wm = [lambda/(n+lambda), 0.5/(n+lambda)*ones(1,2*n)];
    Wc = Wm; Wc(1) = Wc(1) + (1 - alpha^2 + beta);
    Zminus = Z_sig*Wm';
    
    Pzz = R; Pxz = zeros(n,m);
    for j = 1:2*n+1
        Pzz = Pzz + Wc(j)*(Z_sig(:,j)-Zminus)*(Z_sig(:,j)-Zminus)';
        Pxz = Pxz + Wc(j)*(sigma(:,j)-Xminus)*(Z_sig(:,j)-Zminus)';
    end
    
    % 卡尔曼增益
    K = Pxz / Pzz;
    X = Xminus + K*(z - Zminus);
    P = Pminus - K*Pzz*K';
end

function Xnext = state_model(X, Ts, Rr, Xrr, U_qs, U_ds, U_dr, U_qr, T_m, T_e, D, J, C, Tf, I_dl)
    Xnext = X;
    Xnext(1) = X(1) + Ts*((Rr/Xrr)*(U_qs - X(2)) - X(6)*X(1));
    Xnext(2) = X(2) + Ts*((Rr/Xrr)*(U_ds - X(1)) + X(6)*X(2));
    Xnext(3) = X(3) + Ts*(C/Tf*(X(4) - I_dl));
    Xnext(4) = X(4) + Ts*((1/X_dr)*(U_dr - Rr*X(4)));
    Xnext(5) = X(5) + Ts*((1/X_qr)*(U_qr - Rr*X(5)));
    Xnext(6) = X(6) + Ts*((T_m - T_e - D*X(6))/J);
end

function Z = obs_model(X, U_qs, U_ds, Xs, R, omega_s)
    m_d = 0.8*sin(2*pi*50); m_q = 0.5*cos(2*pi*50); % 简化的调制信号
    U_d = m_d*X(3)/2; U_q = m_q*X(3)/2;
    Z = zeros(6,1);
    Z(1) = (omega_s*Xs)/(R+Xs)*X(2) + (R/(R+Xs))*U_q;
    Z(2) = (omega_s*Xs)/(R+Xs)*X(1) - (R/(R+Xs))*U_d;
    Z(3) = X(3);
    Z(4) = (Xs/(R+Xs))*X(3) - (R/(R+Xs))*X(1) + (omega_s*Xs)/(R+Xs)*Z(1);
    Z(5) = (R/(R+Xs))*X(3) + (Xs/(R+Xs))*X(1) + (omega_s*R)/(R+Xs)*Z(2);
    Z(6) = X(6);
end